package easy.trace.core;

import easy.trace.common.Span;
import easy.trace.core.context.TraceContext;
import easy.trace.core.dto.ExceptionType;
import easy.trace.core.dto.TransferParams;
import easy.trace.core.support.RingBuffer;
import easy.trace.core.support.SpanBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;

import java.util.Objects;
import java.util.Optional;

/**
 * @author easy
 */
public class Carrier implements ICarrier {
    private static final Logger log = LoggerFactory.getLogger(Carrier.class);
    private SpanBuilder spanBuilder;
    private static ThreadLocal<TraceContext> contextThreadLocal = new ThreadLocal<>();
    private RingBuffer ringBuffer;

    private static class CarrierHolder {
        private static Carrier INSTANCE = new Carrier();
    }

    public static final Carrier getInstance() {
        return CarrierHolder.INSTANCE;
    }

    public Carrier() {
        this.spanBuilder = SpanBuilder.getSpanContextBuilder();
        this.ringBuffer = new RingBuffer();
        this.ringBuffer.start();
    }

    @Override
    public void init(TraceContext context) {
        Optional.ofNullable(context)
                .ifPresent(con -> {
                    if (con.getServiceId() != null) {
                        contextThreadLocal.set(con);
                    }
                });
    }

    @Override
    public void doBeforeHandle() {
        TraceContext context = contextThreadLocal.get();
        if (Objects.isNull(context)) {
            return;
        }

        try {
            if (context.isCaller()) {
                Span parentSpan = spanBuilder.getParentSpan();
                Span csSpan;
                if (Objects.nonNull(parentSpan)) {
                    csSpan = spanBuilder.createCsSpan(
                            context.getServiceId(),
                            context.getMethod(),
                            context.getEventType(),
                            parentSpan.getTraceId(),
                            parentSpan.getId(),
                            context.address(),
                            context.port());
                } else {
                    csSpan = spanBuilder.createCsSpan(
                            context.getServiceId(),
                            context.getMethod(),
                            context.getEventType(),
                            context.address(),
                            context.port());
                }
                spanBuilder.setCurrentSpan(csSpan);
                context.setTransferParams(csSpan);
            }

            if (context.isCalled()) {
                TransferParams transferParams = context.getTransferParams();
                if (Objects.nonNull(transferParams) && transferParams.getTraceId() != 0L) {
                    Span srSpan = this.spanBuilder.createSrSpan(
                            context.getServiceId(),
                            context.getMethod(),
                            context.getEventType(),
                            transferParams.getTraceId(),
                            transferParams.getParentId(),
                            transferParams.getSelfId(),
                            context.address(),
                            context.port());
                    spanBuilder.setParentSpan(srSpan);
                }
            }
        } catch (Throwable e) {
            log.error(e.getLocalizedMessage());
        }
    }

    @Override
    public void doAfterHandle() {
        TraceContext context = contextThreadLocal.get();
        if (Objects.isNull(context)) {
            return;
        }
        try {
            if (context.isCaller()) {
                Span currentSpan = spanBuilder.getCurrentSpan();
                if (Objects.nonNull(currentSpan)) {
                    spanBuilder.addCrAnnotation(
                            currentSpan,
                            context.address(),
                            context.port());
                    ringBuffer.add(currentSpan);
                }
            }

            if (context.isCalled()) {
                Span parentSpan = spanBuilder.getParentSpan();
                if (Objects.nonNull(parentSpan)) {
                    spanBuilder.addSsAnnotation(
                            parentSpan,
                            context.address(),
                            context.port());
                    ringBuffer.add(parentSpan);
                }
            }
        } catch (Throwable e) {
            log.error(e.getLocalizedMessage());
        } finally {
            contextThreadLocal.remove();
        }
    }

    @Override
    public void doEncounterException(Exception e, ExceptionType et) {
        try {
            final Span parentSpan = spanBuilder.getParentSpan();
            if (Objects.nonNull(parentSpan)) {
                spanBuilder.addExceptionAnnotation(
                        parentSpan,
                        e.getMessage(),
                        et.name());
            }
        } catch (Throwable ex) {
            log.error(ex.getLocalizedMessage());
        }
    }

    @Override
    public void clear() {
        MDC.clear();
    }


}
